package com.huawei.admins.platform.gateway.core.component.limiter;

import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.LoadingCache;
import org.apache.commons.lang3.tuple.Pair;
import org.springframework.util.Assert;

import java.math.BigDecimal;
import java.math.RoundingMode;
import java.time.Duration;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.LongAdder;


public class ServiceRateLimiterTool {

    private final long replenishRate;

    private final long burstCapacity;

    private final long quickThreshold;

    private final int tokenPerRequest;

    private final double fillTime;

    private final long cacheTtl;

    private final LoadingCache<String, TokenPair> tokenLocalCache;

    public ServiceRateLimiterTool(int replenishRate) {
        this(replenishRate, replenishRate, 1);
    }

    public ServiceRateLimiterTool(int replenishRate, int burstCapacity) {
        this(replenishRate, burstCapacity, 1);
    }

    public ServiceRateLimiterTool(int replenishRate, int burstCapacity, int tokenPerRequest) {

        Assert.isTrue(replenishRate > 0, "replenishRate must greater than 0");
        Assert.isTrue(burstCapacity > 0, "burstCapacity must greater than 0");
        Assert.isTrue(tokenPerRequest > 0, "tokenPerRequest must greater than 0");

        this.replenishRate = replenishRate;
        this.burstCapacity = burstCapacity;
        this.tokenPerRequest = tokenPerRequest;
        this.quickThreshold = (long)(burstCapacity * 0.4);
        this.fillTime = new BigDecimal(burstCapacity).divide(new BigDecimal(replenishRate), 1, RoundingMode.HALF_UP).doubleValue();
        this.cacheTtl = (long) Math.floor(fillTime * 2);

        this.tokenLocalCache = Caffeine.newBuilder()
                .expireAfterAccess(Duration.ofSeconds(cacheTtl))
                .build(key -> new TokenPair(key, burstCapacity, 0));
    }

    public double getFillTime() {
        return fillTime;
    }

    public long getCacheTtl() {
        return cacheTtl;
    }

    /*
       下面的快速路径, 实际上还是有并发问题, 因为判断和最后取值并非原子操作.
       但考虑到限流本身不是强精确的业务, 这个并发问题在可接受范围内. 为了换取更高的性能, 不进行处理
     */
    public Pair<Boolean, Long> tryAcquire(String key) {
        TokenPair tokenPair = tokenLocalCache.get(key);
        if (tokenPair == null) {
            return Pair.of(true, burstCapacity);
        }

        // 快速路径，无锁校验
        if (tokenPair.tokenNum.sum() >= quickThreshold + tokenPerRequest) {
            tokenPair.tokenNum.add(-tokenPerRequest);
            return Pair.of(true, tokenPair.tokenNum.sum());
        }

        long now = System.currentTimeMillis();
        long lastRefreshed = tokenPair.lastRefreshTime;
        long delta = Math.max(0, now - lastRefreshed);

        // 满足延迟更新条件：令牌足够且更新时间间隔短 (500ms)
        if (tokenPair.tokenNum.sum() >= tokenPerRequest && delta < 500) {
            tokenPair.tokenNum.add(-tokenPerRequest);
            return Pair.of(true, tokenPair.tokenNum.sum());
        }

        // 进入更新逻辑
        if (delta > 0L) {
            refreshTokenPair(tokenPair, now, delta);
        }
        // 更新完成后再次判断是否拿到令牌
        long currentTokens = tokenPair.tokenNum.sum();
        boolean allowed = currentTokens >= tokenPerRequest;
        if (allowed) {
            tokenPair.tokenNum.add(-tokenPerRequest);
        }
        return Pair.of(allowed, tokenPair.tokenNum.sum());
    }

    private void refreshTokenPair(TokenPair tokenPair, long now, long delta) {
        // 只有一个请求会去更新令牌
        if (tokenPair.updateFlag.compareAndSet(false, true)) {
            try {
                long adjust = Math.max(0, delta * replenishRate / 1000);
                tokenPair.tokenNum.add(adjust);
                tokenPair.lastRefreshTime = now;
            } finally {
                tokenPair.updateFlag.set(false);
            }
        }
        // 其余请求自旋等待更新完成
        int retry = 0;
        while (tokenPair.updateFlag.get() && ++retry < 5) {
            Thread.onSpinWait();
        }
    }

    private static class TokenPair {

        private final String key;

        private final LongAdder tokenNum;

        private final AtomicBoolean updateFlag;

        private long lastRefreshTime;

        public TokenPair(String key, long tokenNum, long lastRefreshTime) {
            this.key = key;
            this.tokenNum = new LongAdder();
            this.updateFlag = new AtomicBoolean(false);
            this.tokenNum.add(tokenNum);
            this.lastRefreshTime = lastRefreshTime;
        }

        @Override
        public String toString() {
            return "TokenPair{" +
                    "key='" + key + '\'' +
                    ", tokenNum=" + tokenNum +
                    ", lastRefreshTime=" + lastRefreshTime +
                    '}';
        }
    }
}
